package com.ruoyi.websocket.interceptor;

import com.ruoyi.framework.web.service.TokenService;
import com.ruoyi.common.core.domain.model.LoginUser;
import com.ruoyi.common.utils.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.net.URI;
import java.util.Map;

/**
 * WebSocket握手拦截器
 * 用于验证用户token和提取用户信息
 * 
 * @author ruoyi
 */
public class WebSocketInterceptor implements HandshakeInterceptor {

    @Autowired
    private TokenService tokenService;

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        
        try {
            // 从URL参数中获取token
            URI uri = request.getURI();
            String query = uri.getQuery();
            String token = null;
            
            if (StringUtils.hasText(query)) {
                String[] params = query.split("&");
                for (String param : params) {
                    if (param.startsWith("token=")) {
                        token = param.substring(6);
                        break;
                    }
                }
            }
            
            if (!StringUtils.hasText(token)) {
                System.out.println("WebSocket连接失败：缺少token参数");
                return false;
            }
            
            // 验证token并获取用户信息
            LoginUser loginUser = tokenService.getLoginUser(token);
            if (loginUser == null) {
                System.out.println("WebSocket连接失败：token无效");
                return false;
            }
            
            // 将用户信息存储到WebSocket会话属性中
            attributes.put("userId", loginUser.getUserId());
            attributes.put("username", loginUser.getUsername());
            attributes.put("loginUser", loginUser);
            
            System.out.println("WebSocket连接成功，用户：" + loginUser.getUsername() + "（ID:" + loginUser.getUserId() + "）");
            return true;
            
        } catch (Exception e) {
            System.err.println("WebSocket握手验证失败：" + e.getMessage());
            return false;
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception exception) {
        if (exception != null) {
            System.err.println("WebSocket握手完成，但有异常：" + exception.getMessage());
        }
    }

}
